# _*_ coding:utf-8 _*_

import numpy as np
import matplotlib.pyplot as plt

from random import shuffle, sample
from math import log2

from utils import properties

np.set_printoptions(suppress=True, linewidth=999, threshold=np.inf)
plt.rc('font', family='Times New Roman')

city_colors = properties.palette('greys')
tour_colors = properties.palette('default')
tour_colors = ['k', 'C0',  'C2', 'C3', 'C4', 'C5', 'C6', 'C7', 'C8', 'C9', 'C1']

class Data:
    def __init__(self, name: str, m: int):
        path = rf'./data/TSPLIB_VC/{name}.tsp/'
        self.coords = np.loadtxt(path + rf'{name}.tsp.txt', skiprows=6, dtype=float)  # read coordinates
        self.__calculate_distances()  # coords distances
        self.cities = self.coords[:, 0].astype(int).tolist()  # city numbers
        self.n = len(self.coords)  # number of cities
        self.A = np.loadtxt(path + rf'access_matrix_sm{m}.txt', skiprows=1)  # read accessibility matrix
        self.m = len(self.A[0]) - 1
        self.depot = np.argwhere(np.sum(self.A[:, 1:], axis=1) == 0)[0][0]  # find depot
        self.S = [[np.where(self.A[i, 1:] == 1)[0]+1][0].tolist() for i, _ in enumerate(self.A)]
        self.S[0] = []  # self.S[i] : accessible salesmen of city i
        self.C = [[np.where(self.A[1:, k] == 1)[0]+1][0].tolist() for k, _ in enumerate(self.A[0])]
        self.C[0] = []  # self.C[i] : accessible cities of salesman i

    def __calculate_distances(self):
        i, x, y = self.coords[:, 0].astype(int), self.coords[:, 1], self.coords[:, 2]
        i = np.insert(i, obj=0, values=0, axis=0)
        x = np.insert(x, obj=0, values=0, axis=0)
        y = np.insert(y, obj=0, values=0, axis=0)
        j = i.reshape(len(i), -1)
        self.distances = np.sqrt((x[i] - x[j]) ** 2 + (y[i] - y[j]) ** 2)
        self.distances[0, :], self.distances[:, 0] = 0, 0

    def cost(self, route=None, routes=None):
        if route is not None:
            dis = 0
            for i in range(1, len(route)):
                dis += self.distances[route[i - 1]][route[i]]
            return dis
        if routes is not None:
            dis = 0
            for k in range(len(routes)):
                for i in range(1, len(routes[k])):
                    dis += self.distances[routes[k][i - 1]][routes[k][i]]
            return dis
        return 0

    def draw(self, plot=plt,  routes=None, title=None,
             rotate=False, emphasize=None, dots=None, dot_size=None, show_text=False):
        dot_size = round(log2(self.n))*6 if dot_size is None else dot_size
        access_num = np.delete(np.sum(self.A[:, 1:], axis=1), obj=0, axis=0)
        coords = np.insert(self.coords[:], obj=3, values=access_num, axis=1)
        if rotate:
            coords[:, 1], coords[:, 2] = coords[:, 2], coords[:, 1]
        access_num = access_num.astype(int)
        plot.scatter(coords[:, 1], coords[:, 2],
                     c=city_colors[coords[:, 3].astype(int)], marker='o', s=dot_size)
        plot.scatter(coords[self.depot-1, 1], coords[self.depot-1, 2], c='k', marker='D', s=dot_size*1.5)
        if dots is not None:
            dots = np.array(dots) - 1
            dots = dots.tolist()
            plot.scatter(coords[dots, 1], coords[dots, 2],
                         c='r', marker='o', s=dot_size)
        # transparent
        if routes is not None:
            alpha = [0.2 for _ in range(len(routes))]
            if emphasize is None:
                emphasize = [e for e in range(len(routes))]
            if type(emphasize) is int:
                emphasize = [emphasize]
            for e in emphasize:
                alpha[e - 1] = 1.0
                emphasize = [emphasize]
        # draw routes
        if routes is not None:
            for i, route in enumerate(routes):
                city_indexes = np.array(route) - 1
                plot.plot(coords[city_indexes, 1], coords[city_indexes, 2], c=tour_colors[i+1],
                          linewidth=2, alpha=alpha[i])
                if show_text:
                    for coord in self.coords[city_indexes]:
                        plot.text(x=coord[1], y=coord[2], s=coord[0].astype(int), alpha=alpha[i])
        # text of numbers
        if show_text and routes is None:
            for coord in self.coords:
                plot.text(x=coord[1], y=coord[2], s=coord[0].astype(int))
        plot.title(title)
        plot.axis('equal')
        return plot

    def init_routes(self, method='greedy'):
        cities = self.cities[:]
        cities.remove(self.depot)
        if method is None:
            print('In init_routes, method not assigned.')
            quit(-1)
        elif method == 'random':
            routes = [[self.depot] for _ in range(self.m)]
            for k in range(self.m):
                c = sample(self.C[k+1], 1)[0]
                cities.remove(c)
                routes[k].append(c)  # add a random city for each route
            shuffle(cities)
            for city in cities:
                s = sample(self.S[city], 1)[0]  # random select an accessible route
                routes[s - 1].append(city)
            for k in range(self.m):
                routes[k].append(self.depot)
            return routes
        elif method == 'greedy':
            routes = [[self.depot] for _ in range(self.m)]
            routes_order = [k for k in range(self.m)]
            while len(cities) > 0:
                shuffle(routes_order)
                for k in routes_order:
                    C = self.C[k+1]
                    city, min_dis = -1, 1e99
                    for c in C:
                        if c in cities:
                            if min_dis > self.distances[routes[k][-1]][c]:
                                min_dis = self.distances[routes[k][-1]][c]
                                city = c
                    if city == -1:
                        routes_order.remove(k)
                    else:
                        routes[k].append(city)
                        cities.remove(city)
            for k in range(self.m):
                routes[k].append(self.depot)
            return routes
        elif method == 'min-insert':
            routes = [[self.depot] for _ in range(self.m)]
            routes_index = list(range(len(routes)))
            shuffle(routes_index)
            for k in routes_index:
                city, min_dis = -1, 1e99
                for c in self.C[k+1]:
                    if c in cities:
                        if min_dis > self.distances[routes[k][-1]][c]:
                            min_dis = self.distances[routes[k][-1]][c]
                            city = c
                routes[k].append(city)
                cities.remove(city)  # add the nearest city for each route
                routes[k].append(self.depot)
            shuffle(cities)
            for city in cities:
                _s, _j, min_cost = -1, -1, 1e99
                for s in self.S[city]:  # search in the route of each salesman
                    route = routes[s - 1]
                    for j in range(1, len(route)):
                        cost = self.distances[route[j - 1]][city] \
                               + self.distances[city][route[j]] - self.distances[route[j - 1]][route[j]]
                        if min_cost > cost:
                            min_cost = cost
                            _s, _j = s, j
                routes[_s - 1].insert(_j, city)
            for route in routes:
                if len(route) < 3:
                    print('Error in route length.')
                    print(routes)
                    quit(-1)
            return routes
        else:
            print(f'In init_routes, method invalid, method={method}.')
            quit(-1)

